/*
 * Copyright 2023-2024, Haiku, Inc. All rights reserved.
 * Distributed under the terms of the MIT License.
 *
 * Authors:
 *		Haiku Development Team
 */

#include <KernelExport.h>

#include <arch/cpu.h>
#include <arch/thread.h>
#include <arch/user_debugger.h>
#include <boot/stage2.h>
#include <kernel.h>
#include <thread.h>
#include <vm/vm_types.h>
#include <vm/VMAddressSpace.h>

#include <string.h>

// #define TRACE_ARCH_THREAD
#ifdef TRACE_ARCH_THREAD
#define TRACE(x) dprintf x
#else
#define TRACE(x) ;
#endif

// 从内核栈中获取iframe
static struct iframe *
get_current_iframe(void)
{
    struct iframe *frame;

    asm volatile("move %0, $sp" : "=r"(frame));

    return frame;
}

// 初始化线程上下文
status_t
arch_thread_init_thread_struct(Thread *thread)
{
    // 为线程分配架构特定的数据结构
    thread->arch_info.fpuEnabled = false;
    thread->arch_info.fpuState = NULL;

    memset(&thread->arch_info.context, 0, sizeof(arch_context));

    return B_OK;
}

// 释放线程上下文
status_t
arch_thread_free_thread_struct(Thread *thread)
{
    // 释放FPU状态
    if (thread->arch_info.fpuState != NULL)
    {
        free(thread->arch_info.fpuState);
        thread->arch_info.fpuState = NULL;
    }

    return B_OK;
}

// 初始化线程局部存储
status_t
arch_thread_init_tls(Thread *thread)
{
    // LoongArch64使用tp寄存器作为线程局部存储指针
    thread->arch_info.context.tp = 0;

    return B_OK;
}

// 初始化内核线程
status_t
arch_thread_init_kernel_thread_stack(Thread *thread, void (*function)(void), void *data)
{
    addr_t *kstack = (addr_t *)thread->kernel_stack_base;
    size_t kstackSize = KERNEL_STACK_SIZE;
    addr_t *stackTop = (addr_t *)(thread->kernel_stack_base + kstackSize);

    // 初始化栈帧
    // 为了简化，我们直接在栈上构建一个iframe结构
    stackTop -= sizeof(iframe) / sizeof(addr_t);
    struct iframe *frame = (struct iframe *)stackTop;

    // 设置程序计数器为函数入口
    frame->era = (addr_t)function;

    // 设置参数寄存器a0为data
    frame->a0 = (addr_t)data;

    // 设置栈指针
    thread->arch_info.context.sp = (addr_t)stackTop;

    // 设置返回地址为thread_kthread_exit
    thread->arch_info.context.ra = (addr_t)thread_kthread_exit;

    return B_OK;
}

// 初始化用户线程
status_t
arch_thread_init_user_thread_stack(Thread *thread, void (*entry)(void), void (*function)(void *), void *data)
{
    // 获取用户栈顶
    addr_t stackTop = thread->user_stack_base + thread->user_stack_size;

    // 确保栈16字节对齐
    stackTop &= ~(addr_t)0xf;

    // 设置初始iframe
    struct iframe frame;
    memset(&frame, 0, sizeof(frame));

    // 设置程序计数器为入口点
    frame.era = (addr_t)entry;

    // 设置参数
    frame.a0 = (addr_t)function;
    frame.a1 = (addr_t)data;

    // 设置栈指针
    frame.sp = stackTop;

    // 将iframe复制到用户栈上
    stackTop -= sizeof(iframe);
    if (user_memcpy((void *)stackTop, &frame, sizeof(frame)) < B_OK)
        return B_BAD_ADDRESS;

    // 设置内核保存的用户栈指针
    thread->arch_info.userFrame = stackTop;

    return B_OK;
}

// 进入用户空间
status_t
arch_thread_enter_userspace(Thread *thread, addr_t entry, addr_t stackFrame, addr_t stackBase, addr_t tlsBase)
{
    struct iframe frame;
    memset(&frame, 0, sizeof(frame));

    // 设置程序计数器为入口点
    frame.era = entry;

    // 设置栈指针
    frame.sp = stackBase + thread->user_stack_size - 16;

    // 设置线程局部存储指针
    frame.tp = tlsBase;

    // 设置参数
    frame.a0 = stackFrame;

    // 进入用户空间
    arch_enter_userspace(&frame);

    // 不应该返回到这里
    panic("arch_thread_enter_userspace() returned!\n");
    return B_ERROR;
}

// 保存FPU状态
static void
save_fpu(Thread *thread, fpu_context *fpuState)
{
    // 保存FPU寄存器
    asm volatile(
        "fst.d $f0, %0, 0\n"
        "fst.d $f1, %0, 8\n"
        "fst.d $f2, %0, 16\n"
        "fst.d $f3, %0, 24\n"
        "fst.d $f4, %0, 32\n"
        "fst.d $f5, %0, 40\n"
        "fst.d $f6, %0, 48\n"
        "fst.d $f7, %0, 56\n"
        "fst.d $f8, %0, 64\n"
        "fst.d $f9, %0, 72\n"
        "fst.d $f10, %0, 80\n"
        "fst.d $f11, %0, 88\n"
        "fst.d $f12, %0, 96\n"
        "fst.d $f13, %0, 104\n"
        "fst.d $f14, %0, 112\n"
        "fst.d $f15, %0, 120\n"
        "fst.d $f16, %0, 128\n"
        "fst.d $f17, %0, 136\n"
        "fst.d $f18, %0, 144\n"
        "fst.d $f19, %0, 152\n"
        "fst.d $f20, %0, 160\n"
        "fst.d $f21, %0, 168\n"
        "fst.d $f22, %0, 176\n"
        "fst.d $f23, %0, 184\n"
        "fst.d $f24, %0, 192\n"
        "fst.d $f25, %0, 200\n"
        "fst.d $f26, %0, 208\n"
        "fst.d $f27, %0, 216\n"
        "fst.d $f28, %0, 224\n"
        "fst.d $f29, %0, 232\n"
        "fst.d $f30, %0, 240\n"
        "fst.d $f31, %0, 248\n"
        : : "r"(fpuState) : "memory");

    // 保存FPU控制寄存器
    uint32 fcsr;
    asm volatile("movfcsr2gr %0, $r0" : "=r"(fcsr));
    fpuState->fcsr = fcsr;
}

// 恢复FPU状态
static void
restore_fpu(Thread *thread, fpu_context *fpuState)
{
    // 恢复FPU寄存器
    asm volatile(
        "fld.d $f0, %0, 0\n"
        "fld.d $f1, %0, 8\n"
        "fld.d $f2, %0, 16\n"
        "fld.d $f3, %0, 24\n"
        "fld.d $f4, %0, 32\n"
        "fld.d $f5, %0, 40\n"
        "fld.d $f6, %0, 48\n"
        "fld.d $f7, %0, 56\n"
        "fld.d $f8, %0, 64\n"
        "fld.d $f9, %0, 72\n"
        "fld.d $f10, %0, 80\n"
        "fld.d $f11, %0, 88\n"
        "fld.d $f12, %0, 96\n"
        "fld.d $f13, %0, 104\n"
        "fld.d $f14, %0, 112\n"
        "fld.d $f15, %0, 120\n"
        "fld.d $f16, %0, 128\n"
        "fld.d $f17, %0, 136\n"
        "fld.d $f18, %0, 144\n"
        "fld.d $f19, %0, 152\n"
        "fld.d $f20, %0, 160\n"
        "fld.d $f21, %0, 168\n"
        "fld.d $f22, %0, 176\n"
        "fld.d $f23, %0, 184\n"
        "fld.d $f24, %0, 192\n"
        "fld.d $f25, %0, 200\n"
        "fld.d $f26, %0, 208\n"
        "fld.d $f27, %0, 216\n"
        "fld.d $f28, %0, 224\n"
        "fld.d $f29, %0, 232\n"
        "fld.d $f30, %0, 240\n"
        "fld.d $f31, %0, 248\n"
        : : "r"(fpuState) : "memory");

    // 恢复FPU控制寄存器
    asm volatile("movgr2fcsr $r0, %0" : : "r"(fpuState->fcsr));
}

// 启用FPU
void arch_thread_context_switch(Thread *from, Thread *to)
{
    // 保存FPU状态（如果启用）
    if (from->arch_info.fpuEnabled)
    {
        if (from->arch_info.fpuState == NULL)
        {
            from->arch_info.fpuState = (fpu_context *)malloc(sizeof(fpu_context));
            if (from->arch_info.fpuState == NULL)
            {
                panic("arch_thread_context_switch: Failed to allocate FPU state!\n");
            }
        }

        save_fpu(from, from->arch_info.fpuState);
    }

    // 切换线程上下文
    switch_threads(from, to);

    // 恢复FPU状态（如果启用）
    if (to->arch_info.fpuEnabled && to->arch_info.fpuState != NULL)
    {
        restore_fpu(to, to->arch_info.fpuState);
    }
}

// 保存线程上下文
void arch_thread_dump_info(void *info)
{
    Thread *thread = (Thread *)info;

    dprintf("thread %" B_PRId32